import os
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
from feature_extract.config import MODEL_CONFIGS
from feature_extract.registry import MODEL_REGISTRY


def preprocess_image(image_path, from_array=False):
    preprocess = transforms.Compose([
        transforms.Resize((256, 128)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    if from_array:
        image = Image.fromarray(image_path)
    else:
        image = Image.open(image_path).convert('RGB')
    image = preprocess(image)
    image = image.unsqueeze(0)  # Add batch dimension
    return image


def extract_features(model, image_path, device, from_array=False, trans=False):
    input_image = preprocess_image(image_path, from_array).to(device)
    if trans:
        with torch.no_grad():
            features = model(input_image, None, mode='featuremap')
            return features.squeeze().cpu().numpy()
    else:
        with torch.no_grad():
            features = model(input_image)
        return features.squeeze().cpu().numpy()


import torch
from collections import OrderedDict


def remove_prefix(state_dict, prefix):
    """Old style model is stored with all names of parameters sharing common prefix 'module.'"""
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        if k.startswith(prefix):
            new_key = k[len(prefix):]
        else:
            new_key = k
        new_state_dict[new_key] = v
    return new_state_dict


def add_prefix(state_dict, prefix):
    new_state_dict = {}
    for key in state_dict.keys():
        new_key = prefix + key
        new_state_dict[new_key] = state_dict[key]
    return new_state_dict


def replace_suffix(state_dict, suffix_to_remove, suffix_to_add):
    """替换键名中的后缀"""
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        if suffix_to_remove in k:
            new_key = k.replace(suffix_to_remove, suffix_to_add)
        else:
            new_key = k
        new_state_dict[new_key] = v
    return new_state_dict


def remove_keys(state_dict, keys):
    """移除指定键名"""
    for key in keys:
        if key in state_dict:
            del state_dict[key]
    return state_dict


def load_classifier_weights(classifier_path):
    """加载 classifier 层的权重"""
    classifier_checkpoint = torch.load(classifier_path, encoding='iso-8859-1')
    classifier_state_dict = {
        'classifier.weight': classifier_checkpoint['state_dict']['classifier.weight'],
        'classifier.bias': classifier_checkpoint['state_dict']['classifier.bias'],
    }
    return classifier_state_dict


def cosine_similarity(vector1, vector2, trans=False, model=None, device='cuda'):
    if not trans:
        dot_product = np.dot(vector1, vector2)
        norm1 = np.linalg.norm(vector1)
        norm2 = np.linalg.norm(vector2)
        return dot_product / (norm1 * norm2)
    else:
        with torch.no_grad():
            similarity= model(torch.from_numpy(vector1).float().unsqueeze(0).to(device),
                                     torch.from_numpy(vector2).float().unsqueeze(0).to(device), mode='compare')
        similarity = torch.sigmoid(similarity).cpu().numpy().squeeze().squeeze()
        return similarity


def initialize_model(model_name: str, device=torch.device('cpu')):
    print(MODEL_REGISTRY._registry.keys())
    model_class = MODEL_REGISTRY.get(model_name)
    model_params = MODEL_CONFIGS[model_name]['params']
    model = model_class(**model_params)
    method_dict = MODEL_CONFIGS[model_name]['load_weight_method']
    method_dict['path'] = MODEL_CONFIGS[model_name]['weights']
    method_dict['model'] = model
    method_dict['device'] = device
    method = MODEL_REGISTRY.get(method_dict['name'])
    model = method(**method_dict)
    model.to(device)
    return model


def initialize_train_model(model_name: str, device=torch.device('cpu')):
    model_class = MODEL_REGISTRY.get(model_name)
    model_params = MODEL_CONFIGS[model_name]['params']
    model = model_class(**model_params)
    model.to(device)
    return model


def find_most_similar_class(input_features, class_features):
    max_similarity = -1
    most_similar_class = None
    resdict = {}
    for class_name, avg_features in class_features.items():
        similarity = cosine_similarity(input_features, avg_features)  # 计算余弦相似度
        resdict[class_name] = similarity
        if similarity > max_similarity:
            max_similarity = similarity
            most_similar_class = class_name
    return most_similar_class, max_similarity, resdict


def find_most_similar_class_merge(input_features, class_features, trans=False, model=None, device='cuda'):
    max_similarity = -1
    most_similar_class = None
    resdict = {}
    for class_name, feature_list in class_features.items():
        if not trans:
            #print('not trans')
            similarity = cosine_similarity(input_features, feature_list[0])  # 计算余弦相似度
        else:
            with torch.no_grad():
                similarity= model(torch.from_numpy(input_features).float().unsqueeze(0).to(device),
                                         torch.from_numpy(feature_list[0]).float().unsqueeze(0).to(device), mode='compare')
            similarity = torch.sigmoid(similarity).cpu().numpy().squeeze().squeeze()
        resdict[class_name] = similarity
        if similarity > max_similarity:
            max_similarity = similarity
            most_similar_class = class_name

    return most_similar_class, max_similarity, resdict


def get_all_jpg_files(directory):
    jpg_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.jpg'):
                jpg_files.append(os.path.join(root, file))
    return jpg_files
